from torch.utils.data import Dataset
from os.path import join as o_p_join
from PIL import Image
from torchvision.transforms import Compose, ToTensor, Normalize


NUM_CLASSES = 40
DL_BSIZE = 512


class CFP_Dset(Dataset):
    def __init__(self, root_dir, train=True):
        super(CFP_Dset, self).__init__()
        # 创建类型列表
        self.classes = []
        for i in range(NUM_CLASSES):
            self.classes.append('person-' + str(i + 1))
        # 统计各项数据列表
        self.root_dir = root_dir
        self.item_dir = []
        self.item_class = []
        self.item_count = 0
        for i in range(NUM_CLASSES):
            for j in range(10):
                if (train and (j == 9)) or ((not train) and (j != 9)):
                    continue
                self.item_count += 1
                self.item_dir.append(str(10 * i + j + 1) + '.bmp')
                self.item_class.append(i)
        # 确定数据预处理方式
        self.item_trans = Compose([ToTensor(), Normalize(0.5, 0.2)])

    def __getitem__(self, idx):
        img_dir = o_p_join(self.root_dir, self.item_dir[idx])
        img = self.item_trans(Image.open(img_dir))
        return img, self.item_class[idx]

    def __len__(self):
        return self.item_count
